#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
密钥生成脚本
用于生成RSA密钥对和AES密钥
"""

import os
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.serialization import Encoding, PrivateFormat, PublicFormat, NoEncryption
from cryptography.fernet import Fernet
import secrets

def generate_rsa_keys():
    """
    生成RSA密钥对
    """
    print("正在生成RSA密钥对...")
    
    # 生成私钥
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
    )
    
    # 获取公钥
    public_key = private_key.public_key()
    
    # 序列化私钥
    private_pem = private_key.private_bytes(
        encoding=Encoding.PEM,
        format=PrivateFormat.PKCS8,
        encryption_algorithm=NoEncryption()
    )
    
    # 序列化公钥
    public_pem = public_key.public_bytes(
        encoding=Encoding.PEM,
        format=PublicFormat.SubjectPublicKeyInfo
    )
    
    return private_pem, public_pem

def save_keys_to_files(private_key, public_key):
    """
    保存密钥到文件
    """
    # 创建keys目录
    if not os.path.exists('keys'):
        os.makedirs('keys')
    
    # 保存私钥
    with open('keys/server_private_key.pem', 'wb') as f:
        f.write(private_key)
    
    # 保存公钥
    with open('keys/server_public_key.pem', 'wb') as f:
        f.write(public_key)
    
    print("RSA密钥已保存到 keys/ 目录")

def generate_aes_key():
    """
    生成AES密钥
    """
    print("正在生成AES密钥...")
    key = Fernet.generate_key()
    
    # 创建keys目录
    if not os.path.exists('keys'):
        os.makedirs('keys')
    
    # 保存AES密钥
    with open('keys/aes_key.key', 'wb') as f:
        f.write(key)
    
    print("AES密钥已保存到 keys/aes_key.key")
    return key

def main():
    print("=== 密钥生成工具 ===")
    print("正在生成加密通信所需的密钥...")
    
    # 生成RSA密钥对
    private_key, public_key = generate_rsa_keys()
    save_keys_to_files(private_key, public_key)
    
    # 生成AES密钥
    aes_key = generate_aes_key()
    
    print("\n密钥生成完成！")
    print("生成的文件：")
    print("- keys/server_private_key.pem (服务端RSA私钥)")
    print("- keys/server_public_key.pem (服务端RSA公钥)")
    print("- keys/aes_key.key (AES对称密钥)")
    print("\n请妥善保管这些密钥文件！")

if __name__ == "__main__":
    main()